import os
import json
from PIL import Image
from tqdm import tqdm

# InfantImageDataset
class CropPadding:
    def __init__(self, box=(80, 0, 1570, 1200)):
        self.box = box

    def __call__(self, img):
        return img.crop(self.box)

def main(data_path, split_name, tar_dir):
    with open(os.path.join(data_path, 'annotations.json')) as f:
        data_dict = json.load(f)
    
    with open(os.path.join(data_path, 'split', f'{split_name}.json')) as f:
        split_list = json.load(f)['test']
    
    os.makedirs(os.path.join(data_path, tar_dir), exist_ok=True)
    
    for image_name in tqdm(split_list, desc="Processing images"):
        data = data_dict[image_name]
        img_path = os.path.join(data_path, data['image_path'])
        img = Image.open(img_path)
        img = CropPadding()(img)
        img = img.resize((224, 224))
        
        tar_path = os.path.join(data_path, tar_dir, image_name)
        data_dict[image_name]['resize_image_path'] = tar_path
        
        img.save(tar_path)
    
    with open(os.path.join(data_path, 'annotations.json'), 'w') as f:
        json.dump(data_dict, f)
    
if __name__ == '__main__':
    data_path = '/mnt/d/Workspace/MiniReti/'
    split_name = 'all'
    tar_dir = 'resize_images'
    main(data_path, split_name, tar_dir)
